package com.xr.core.system.core.aop;

import cn.hutool.core.util.StrUtil;
import com.xr.core.system.common.annotation.TrafficLimit;
import com.xr.core.system.common.enums.LimitTypeEnum;
import com.xr.core.system.common.exception.ServiceException;
import com.xr.core.system.core.utils.IpUtil;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.Collections;

@Aspect
@Component
@Slf4j
public class LimitAspect {

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    @Pointcut("@annotation(trafficLimit)")
    public void limitPointCut(TrafficLimit trafficLimit) {
    }

    @Around(value = "limitPointCut(trafficLimit)", argNames = "pjp,trafficLimit")
    public Object interceptor(ProceedingJoinPoint pjp, TrafficLimit trafficLimit) {
        MethodSignature signature = (MethodSignature) pjp.getSignature();
        Method method = signature.getMethod();
        LimitTypeEnum limitType = trafficLimit.limitType();
        String key;
        int limitPeriod = trafficLimit.period();
        int limitCount = trafficLimit.count();
        // 根据限流类型获取不同的key ,如果不传我们会以方法名作为key
        switch (limitType) {
            case IP:
                key = IpUtil.getClientIp();
                break;
            case CUSTOMER:
                key = trafficLimit.key();
                break;
            default:
                key = method.getName();
        }
        try {
            String luaScript = buildLuaScript();
            RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
            Number count = redisTemplate.execute(redisScript, Collections.singletonList(StrUtil.join(trafficLimit.prefix(), key)), limitCount, limitPeriod);
            log.info("Access try count is {}  and key = {}", count, key);
            if (count != null && count.intValue() <= limitCount) {
                return pjp.proceed();
            } else {
                throw new ServiceException("服务器繁忙，请稍后再试！");
            }
        } catch (Throwable e) {
            if (e instanceof RuntimeException) {
                throw new RuntimeException(e.getLocalizedMessage());
            }
            throw new RuntimeException("server exception");
        }
    }

    /**
     * 编写 redis Lua 限流脚本
     */
    public String buildLuaScript() {
        return "local c" +
                "\nc = redis.call('get',KEYS[1])" +
                // 调用不超过最大值，则直接返回
                "\nif c and tonumber(c) > tonumber(ARGV[1]) then" +
                "\nreturn c;" +
                "\nend" +
                // 执行计算器自加
                "\nc = redis.call('incr',KEYS[1])" +
                // 从第一次调用开始限流，设置对应键值的过期
                "\nif tonumber(c) == 1 then" +
                "\nredis.call('expire',KEYS[1],ARGV[2])" +
                "\nend" +
                "\nreturn c;";
    }

}
